import logging
import json


from ..dantic import (
    ParamOutline,
    ParamTitle,
    ParamGetContentFromOutline,
    ParamGenerateSingleParagraph,
    ParamRegenDigest,
)
from .base import BaseAiModel, purgeJsonPrefix, InvokeResult
from . import prompt_factory


from common import ai


LOGGER = logging.getLogger(__name__)


class ModelProxy(BaseAiModel):

    def __init__(self, modelName: str) -> None:
        super().__init__()
        self.modelName = modelName
        # 这个参数是摆设，后边都是qianwen，这个参数会显示在页面上。写死把。
        self.modelName = 'qwen-max'
        LOGGER.info(f"ModelProxy 实例化.模型:{modelName}")

    def _sendRequest(self, userMessage: str) -> InvokeResult:
        try:
            aiModel = ai.getAdapter(ai.MODEL_QIANWEN)
            queryResult: ai.AiQueryResult = aiModel.query(
                ai.AiQuery(userMessage=userMessage)
            )

            if not queryResult.isValid:
                return InvokeResult(isSuccess=False, message=queryResult.errMessage)

            return InvokeResult(
                isSuccess=True, data=queryResult.text, tokenSize=queryResult.totalToken
            )
        except Exception as e:
            LOGGER.exception(e)
            return InvokeResult(isSuccess=False, message=str(e))

    def getTitle(self, param: ParamTitle) -> InvokeResult:

        result: InvokeResult = self._sendRequest(prompt_factory.getTitle(param))
        if result.isSuccess:
            result.data = purgeJsonPrefix(result.data)
        return result

    def getOutline(self, param: ParamOutline) -> InvokeResult:
        # 这里需要考虑调用三次GPT，把三次返回的提纲按照格式，一次性的放到返回文件中发送到前端。
        finalResult = []
        finalTokenSize = 0
        thesis_styles = [
            {"author": "行业专家", "style": "加入更多的对该论题的深度理解和前沿内容"},
            {
                "author": "大学学生",
                "style": "从基础理论入手，多分析现有数据，展示自己的学习成果",
            },
        ]

        for style in thesis_styles:
            try:

                param.author = style["author"]
                param.style = style["style"]
                user_message = prompt_factory.getOutline(param)

                result: InvokeResult = self._sendRequest(user_message)
                if result.isSuccess:
                    resultData = purgeJsonPrefix(str(result.data))

                    finalResult.append(json.loads(resultData))
                    finalTokenSize += result.tokenSize

                # debug
            except Exception as e:
                LOGGER.exception(e)

        return InvokeResult(data=finalResult, tokenSize=finalTokenSize)
    
    def generateParagraph(self, param: ParamGenerateSingleParagraph) -> InvokeResult:
        result: InvokeResult = self._sendRequest(
            prompt_factory.generateSingleParagraph(param)
        )
        if result.isSuccess:
            result.data = purgeJsonPrefix(result.data)
        return result

    def generateDigest(self, param: ParamRegenDigest) -> InvokeResult:

        result: InvokeResult = self._sendRequest(prompt_factory.generateDigest(param))
        if result.isSuccess:
            resBody = purgeJsonPrefix(result.data)
            resBody = resBody.replace("\n", "")
            resBody = resBody.replace("\r", "")
            try:
                resultJson = json.loads(resBody)
                keywords = resultJson["keywords"]
                if keywords is not None and type(keywords) == list:
                    resultJson["keywords"] = ",".join(keywords)
                result.data = resultJson
            except Exception as e:
                LOGGER.exception(e)
        return result

    def generateDigestEnglish(self, param: ParamRegenDigest) -> InvokeResult:

        result: InvokeResult = self._sendRequest(
            prompt_factory.generateDigestEnglish(param)
        )
        if result.isSuccess:
            resBody = purgeJsonPrefix(result.data)
            resBody = resBody.replace("\n", "")
            resBody = resBody.replace("\r", "")
            try:
                # userMessage 里要求输出 digest的json额结构，在这里解除一下结构
                resultJson = json.loads(resBody)
                # 如果关键词给出的是list，变成str，
                keywords = resultJson["keywords"]
                if keywords is not None and type(keywords) == list:
                    resultJson["keywords"] = ",".join(keywords)
                result.data = resultJson
            except Exception as e:
                LOGGER.exception(e)
        return result


